########################################################################
#    Create model script used by WinBUGS for fitting  
#    dynamic linear mixed models (dlmm)
# 
#    Author: Wayne Zhang, June 2012 
#            actuary_zhang@hotmail.com
########################################################################

# get family dist'n and link (adapted from glmmBUGS)
dist_link <- function(family){
  if (family == "poisson"){
    link <- "log"
    ddist <- "dpois"
  } else if (family == "bernoulli"){
    link <- "logit"
    ddist <- "dbern"
  } else if (family == "binomial"){
    link <- "logit"
    ddist <- "dbin"
  } else if (family == "gaussian"){
    link <- ""
    ddist <- "dnorm"
  } else {
    ddist <- paste0("d", family)
    link <- ""
  }
  if (link != "") {
    link <- paste0(link, "(")
    endlink <- ")"
  } else {
    endlink <- ""
  }
  list(ddist = ddist, link = link, endlink = endlink)
}

# concatenate the mean specifications from each piece
link_mu <- function(lst, collapse = " +\n"){
  lst <- lst[which(sapply(lst, length) > 0)]
  paste0(lapply(lst, "[[", "mu"), collapse = collapse)
}

# concatenate the prior distributions (prior + hyperprior)
link_prior <- function(lst){
  lst <- lst[which(sapply(lst, length) > 0)]
  paste0(lapply(lst, "[[", "prior"), collapse = "\n")
}

# prior specification for covariance matrix
model_Sigma <- function(name, d = 2, C = 1){
  # specify the distribution of Sigma
  ind <- paste0("[", "1:", d, ", 1:", d, "]")
  Sigma <- paste0("tau.", name, ind, " ~ dwish(prec.", name, "[ , ], ", d, ")\n", 
                  "sigma.", name, ind, " <- inverse(tau.", name, "[ , ])\n")
  
  # prior precision
  prec <- c()  
  for (i in seq_len(d)){
    for (j in seq_len(d)){
      prec <- c(prec, paste0("prec.", name, "[", i, ", ", j, "] <- ", 
                             ifelse(i == j, C, 0), "\n"))
    }
  }
  prec <- paste(prec, collapse = "")
  paste0(Sigma, prec)
}

# prior specification for variance
model_sigma <- function(name, C = 0.001){
  paste0("tau.", name, " ~ dgamma(", C, ", ", C, ")\n",
         "sigma.", name, " <- 1/sqrt(tau.", name, ")\n")
}

# model specification for fixed effects
model_fixed <- function(p){
  if (p == 0){
    return(NULL)
  } else {      
    mu <- paste("X[i, ", 1:p, "] * beta[", 1:p, ", 1]", sep = "", collapse = " + ")
    prior <- paste("for (i in 1:", p, "){\n",
      "  beta[i, 1] ~ dnorm(0, 1.0E-6)\n", 
      "}\n", sep = "")
    return(list(mu = mu, prior = prior))
  }
}

# model specification for random effects
model_random <- function(nt, nc, nlev){
  if (nt == 0){
    return(NULL)
  } else {
    script <- lapply(1:nt, function(x){
      Z <- paste0("Z", x)
      b <- paste0("b", x)
      grp <- paste0("grp", x)
      mu <- paste0(Z, "[i, ", 1:nc[x], "] * ", b, "[", grp, "[i], ", 1:nc[x], "]",
                   collapse = " + ")
      if (nc[x] > 1){
        prior <- paste0("for (i in 1:", nlev[x], "){\n",
                       "  ", b, "[i, 1:", nc[x], "] ~ dmnorm(mu.", b, "[], tau.", b, "[ , ])\n", 
                       "}\n", 
                        paste0("mu.", b, "[", 1:nc[x], "] <- 0", collapse = "\n"), "\n\n",
                        model_Sigma(b, nc[x]))
      } else{    
        prior <- paste0("for (i in 1:", nlev[x], "){\n",
                        "  ", b, "[i, 1] ~ dnorm(0, tau.", b, ")\n", 
                        "}\n\n",
                        model_sigma(b))                                   
      }
      list(mu = mu, prior = prior)
    })
    # combine all random effects 
    mu <- link_mu(script, " + ")
    prior <- link_prior(script)
    return(list(mu = mu, prior = prior))
  }
}

# model specification for dynamic effects
model_dynamic <- function(r, t){
  mu <- paste0("F[i, ", 1:r, "] * theta[time[i], ", 1:r, "]", collapse = " + ")  
  if (r > 1){
    ind <- paste0("[1:", r, ", ", "1:", r, "]")    
    prior <- paste0("for (i in 1:", r, "){\n",
                    "  theta[1, i] ~ dnorm(0, 1.0E-6)\n",
                    "}\n\n",                
                    "for (i in 2:", t, "){\n",
                    "  theta[i, 1:", r, "] ~ dmnorm(theta[i - 1, ], tau.theta[ , ])\n",
                    "}\n\n", 
                    model_Sigma("theta", r))
  } else{  
    prior <- paste0("theta[1, 1] ~ dnorm(0, 1.0E-6)\n\n",    
                    "for (i in 2:", t, "){\n",
                    "  theta[i, 1] ~ dnorm(theta[i - 1, 1], tau.theta)\n",
                    "}\n\n", 
                    model_sigma("theta"))
  }
  list(mu = mu, prior = prior)
}

# function to generate the model script to "model.file"
bugs_model <- function(family, dims, nc, nlev, model.file){
  fixed <- model_fixed(dims["n.beta"])
  random <- model_random(dims["n.term"], nc, nlev)
  dynamic <- model_dynamic(dims["n.theta"], dims["n.t"])
  mu <- link_mu(list(fixed, random, dynamic), collapse = " +\n      ")
  prior <- link_prior(list(fixed, random, dynamic))
  
  # writing script to file
  sink(file = model.file)
  cat("model{\n\n")
  cat(paste0("for (i in 1:", dims["n.obs"], "){\n"))
  
  # distribution for data
  fam <- dist_link(family)
  disty <- paste0("  Y[i] ~ ", fam$ddist, "(mu[i]")
  if (family == "gaussian"){
    disty <- paste0(disty, ", ", ifelse(dims["n.wts"], "wts[i] *", ""), "tau.y)\n")
  } else{
    disty <- paste0(disty, ")\n")
  }
  cat(disty)
  
  # mean equation
  cat(paste0("  ", fam$link, "mu[i]", fam$endlink, " <- ", mu, 
         ifelse(dims["n.off"], " + off[i]\n", "\n")))
  cat("}\n\n") 
  
  # observation variation for gaussian
  if (family == "gaussian") 
    cat(model_sigma("y"))
  cat("\n")
  
  # parameter distributions
  cat(prior)
  cat("\n}")
  
  # finish writing
  sink()  
}



